{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/mnt/tess/astronet/checkpoints/fa1_38_run_1/1/AstroCNNModel_final_alpha_1_20220504_164445',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/2/AstroCNNModel_final_alpha_1_20220504_172032',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/3/AstroCNNModel_final_alpha_1_20220504_175419',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/4/AstroCNNModel_final_alpha_1_20220504_182735',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/5/AstroCNNModel_final_alpha_1_20220504_190055',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/6/AstroCNNModel_final_alpha_1_20220504_193432',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/7/AstroCNNModel_final_alpha_1_20220504_200824',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/8/AstroCNNModel_final_alpha_1_20220504_204203',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/9/AstroCNNModel_final_alpha_1_20220504_211537',\n",
       " '/mnt/tess/astronet/checkpoints/fa1_38_run_1/10/AstroCNNModel_final_alpha_1_20220504_214944']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "chkpt_root = '/mnt/tess/astronet/checkpoints/fa1_38_run_1'\n",
    "data_files = '/mnt/tess/astronet/tfrecords-38-train/*'\n",
    "tces_file = '/mnt/tess/astronet/tces-v14-train.csv'\n",
    "\n",
    "nruns = 10\n",
    "\n",
    "def load_ensemble(chkpt_root, nruns):\n",
    "    checkpts = []\n",
    "    for i in range(nruns):\n",
    "        parent = os.path.join(chkpt_root, str(i + 1))\n",
    "        if not os.path.exists(parent):\n",
    "            break\n",
    "        all_dirs = os.listdir(parent)\n",
    "        if not all_dirs:\n",
    "            break\n",
    "        d, = all_dirs\n",
    "        checkpts.append(os.path.join(parent, d))\n",
    "    return checkpts\n",
    "\n",
    "paths = load_ensemble(chkpt_root, nruns)\n",
    "paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running model 1\n",
      "19919 records\n",
      "Running model 2\n",
      "19919 records\n",
      "Running model 3\n",
      "19919 records\n",
      "Running model 4\n",
      "19919 records\n",
      "Running model 5\n",
      "19919 records\n",
      "Running model 6\n",
      "19919 records\n",
      "Running model 7\n",
      "19919 records\n",
      "Running model 8\n",
      "19919 records\n",
      "Running model 9\n",
      "19919 records\n",
      "Running model 10\n",
      "19919 records\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "from astronet import predict\n",
    "import tensorflow as tf\n",
    "\n",
    "\n",
    "def run_predictions(path):\n",
    "    predict.FLAGS = predict.parser.parse_args([\n",
    "      '--model_dir', path,\n",
    "      '--data_files', data_files,\n",
    "      '--output_file', '',\n",
    "    ])\n",
    "\n",
    "    return predict.predict()\n",
    "\n",
    "\n",
    "paths = load_ensemble(chkpt_root, nruns)\n",
    "ensemble_preds = []\n",
    "config = None\n",
    "for i, path in enumerate(paths):\n",
    "    print(f'Running model {i + 1}')\n",
    "    preds, config = run_predictions(path)\n",
    "    ensemble_preds.append(preds.set_index('astro_id'))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = ['disp_e', 'disp_n', 'disp_j', 'disp_s', 'disp_b']\n",
    "\n",
    "col_e = labels.index('disp_e')\n",
    "thresh = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "agg_preds = {}\n",
    "\n",
    "for preds in ensemble_preds:\n",
    "    for ex_id in preds.index:\n",
    "        if ex_id not in agg_preds:\n",
    "            agg_preds[ex_id] = []\n",
    "\n",
    "        row = preds[preds.index == ex_id]\n",
    "        pred_v = row.values[0]\n",
    "        if len(row.values) > 1:\n",
    "            print(f'Warning: duplicate predictions for {ex_id}')\n",
    "        if pred_v[col_e] >= thresh:\n",
    "            agg_preds[ex_id].append('disp_e')\n",
    "        else:\n",
    "            masked_v = [v if i != col_e else 0 for i, v in enumerate(pred_v)]\n",
    "            agg_preds[ex_id].append(preds.columns[np.argmax(masked_v)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_preds = []\n",
    "for ex_id in list(agg_preds.keys()):\n",
    "    counts = {l: 0 for l in labels}\n",
    "    for e in agg_preds[ex_id]:\n",
    "        counts[e] += 1\n",
    "    maxcount = max(counts.values())\n",
    "    counts.update({\n",
    "        'astro_id': ex_id,\n",
    "        'maxcount': maxcount,\n",
    "    })\n",
    "    final_preds.append(counts)\n",
    "    \n",
    "final_preds = pd.DataFrame(final_preds).set_index('astro_id')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tce_table = pd.read_csv(tces_file, header=0, low_memory=False)\n",
    "tce_table['astro_id'] = tce_table['Astro ID']\n",
    "tce_table = tce_table.set_index('astro_id')\n",
    "for l in labels:\n",
    "    tce_table[l] = tce_table[l[:-1] + l[-1].upper()]\n",
    "tce_labels = tce_table[labels + ['TIC ID']]\n",
    "\n",
    "pl = final_preds.join(tce_labels, on='astro_id', how='left', lsuffix='_p')\n",
    "\n",
    "pl.head()\n",
    "pd.set_option('display.max_columns', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall: 0.9952426260704091\n",
      "Precision: 0.8\n"
     ]
    }
   ],
   "source": [
    "ppos = (pl['disp_e_p'] > 0)\n",
    "pos = (pl['disp_e'] > 0)\n",
    "\n",
    "pneg = (pl['disp_e_p'] == 0)\n",
    "neg = (pl['disp_e'] == 0)\n",
    "\n",
    "print('Recall:', len(pl[ppos & pos]) / len(pl[pos]))\n",
    "print('Precision:', len(pl[ppos & pos]) / len(pl[ppos]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "49799681\n",
      "50905927\n",
      "761960972\n",
      "379464439\n",
      "369860950\n",
      "376936788\n",
      "381366555\n",
      "176582931\n",
      "383716793\n",
      "387829772\n"
     ]
    }
   ],
   "source": [
    "for i in pl[pos & pneg]['TIC ID']:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30313096\n",
      "231792014\n",
      "469269031\n",
      "404851658\n",
      "50269985\n",
      "339959047\n",
      "31415158\n",
      "89307726\n",
      "279614421\n",
      "384077498\n",
      "149928367\n",
      "260700439\n",
      "186302615\n",
      "424277160\n",
      "457139941\n",
      "173008159\n",
      "404355884\n",
      "252809234\n",
      "1884267302\n",
      "235562906\n",
      "154006988\n",
      "1718006824\n",
      "105378745\n",
      "252733538\n",
      "240341734\n",
      "446012335\n",
      "308958576\n",
      "332867524\n",
      "170932338\n",
      "372186057\n",
      "98937256\n",
      "85517073\n",
      "388906923\n",
      "262169297\n",
      "154699047\n",
      "445417446\n",
      "83408987\n",
      "356822872\n",
      "99382119\n",
      "380289423\n",
      "43667308\n",
      "343463316\n",
      "190622748\n",
      "405460905\n",
      "322399290\n",
      "277298580\n",
      "404931649\n",
      "306890368\n",
      "257772292\n",
      "410228822\n",
      "425083264\n",
      "284196025\n",
      "86144938\n",
      "88353444\n",
      "349572054\n",
      "201563776\n",
      "396967912\n",
      "31391041\n",
      "340217291\n",
      "300557619\n",
      "279615956\n",
      "276793626\n",
      "373843857\n",
      "301955257\n",
      "409639832\n",
      "261136311\n",
      "40515061\n",
      "90161731\n",
      "30641264\n",
      "363414722\n",
      "312028996\n",
      "360961597\n",
      "363979429\n",
      "271597452\n",
      "327595687\n",
      "101821550\n",
      "86516107\n",
      "233896030\n",
      "148392230\n",
      "198536950\n",
      "514546468\n",
      "284701173\n",
      "60385535\n",
      "405864709\n",
      "32000625\n",
      "254158487\n",
      "351472293\n",
      "30267088\n",
      "384195094\n",
      "295372927\n",
      "396957480\n",
      "389363745\n",
      "270365850\n",
      "323170106\n",
      "349635294\n",
      "420174832\n",
      "164464417\n",
      "239551545\n",
      "92328347\n",
      "417057112\n",
      "1715180951\n",
      "468371813\n",
      "378275445\n",
      "277566483\n",
      "452741460\n",
      "82308728\n",
      "239209638\n",
      "64187180\n",
      "329704753\n",
      "428203967\n",
      "352317188\n",
      "273366789\n",
      "260304296\n",
      "293507357\n",
      "394233079\n",
      "277297046\n",
      "271577058\n",
      "382432832\n",
      "406240476\n",
      "370227480\n",
      "270312314\n",
      "343578197\n",
      "279958091\n",
      "55660976\n",
      "455698002\n",
      "156993882\n",
      "333064504\n",
      "286631521\n",
      "359729929\n",
      "235995157\n",
      "121741238\n",
      "92774987\n",
      "354624921\n",
      "387445439\n",
      "420177694\n",
      "244857422\n",
      "12675678\n",
      "137147023\n",
      "285792766\n",
      "459218950\n",
      "348468490\n",
      "138576706\n",
      "66727473\n",
      "470358759\n",
      "273662394\n",
      "335951920\n",
      "141610473\n",
      "404851347\n",
      "451540413\n",
      "293438120\n",
      "391745260\n",
      "404266600\n",
      "31634175\n",
      "87339522\n",
      "269451168\n",
      "295253049\n",
      "40342199\n",
      "384823413\n",
      "220375993\n",
      "326693703\n",
      "380457228\n",
      "79655149\n",
      "279161495\n",
      "40974499\n",
      "30031180\n",
      "276553165\n",
      "122626372\n",
      "257582255\n",
      "233097244\n",
      "239294743\n",
      "267694283\n",
      "316413911\n",
      "154068411\n",
      "122220263\n",
      "137906810\n",
      "116113749\n",
      "229786214\n",
      "193615129\n",
      "258827116\n",
      "357036318\n",
      "256298025\n",
      "282173645\n",
      "161563111\n",
      "298734307\n",
      "459976751\n",
      "122280962\n",
      "341082784\n",
      "201731557\n",
      "23497144\n",
      "179308949\n",
      "309654719\n",
      "293464697\n",
      "273514654\n",
      "310360141\n",
      "21995336\n",
      "182450724\n",
      "277022838\n",
      "90716944\n",
      "31961394\n",
      "198051336\n",
      "75687121\n",
      "261422404\n",
      "238229126\n",
      "389668544\n",
      "200230834\n",
      "257772293\n",
      "407063129\n",
      "404603404\n",
      "404768904\n",
      "382625239\n",
      "276796826\n",
      "404035038\n",
      "277022748\n",
      "224692180\n",
      "258351350\n",
      "139124331\n",
      "173710176\n",
      "239814191\n",
      "82620914\n",
      "266908723\n",
      "362259855\n",
      "101232323\n",
      "105450460\n",
      "10935045\n",
      "457681334\n",
      "458108252\n",
      "463974233\n",
      "47390165\n",
      "48780110\n",
      "4933576\n",
      "50269985\n",
      "60647894\n",
      "62971651\n",
      "67688551\n",
      "76698707\n",
      "82806895\n",
      "90128622\n",
      "98548087\n",
      "363485837\n",
      "62023270\n",
      "380782608\n",
      "453100472\n",
      "363982405\n",
      "112716272\n",
      "389129453\n",
      "307155995\n",
      "30113151\n",
      "373523051\n",
      "23722025\n",
      "370109164\n",
      "456556422\n",
      "381798773\n",
      "260414617\n",
      "382626168\n",
      "456611584\n",
      "220371571\n",
      "287775969\n",
      "349573707\n",
      "349373192\n",
      "128882149\n",
      "120401906\n",
      "349485058\n",
      "384608233\n",
      "370101492\n",
      "349831164\n",
      "74649892\n",
      "275761018\n",
      "63931474\n",
      "445181581\n",
      "418749263\n",
      "431380283\n",
      "180416902\n",
      "238893940\n",
      "166280254\n",
      "294954838\n",
      "179305059\n",
      "394372611\n",
      "456516439\n",
      "87168122\n",
      "326606721\n",
      "404851226\n",
      "277299710\n",
      "230968587\n",
      "29984021\n",
      "369932580\n",
      "87164255\n",
      "306352399\n",
      "394345045\n",
      "177308777\n",
      "411917726\n",
      "294092960\n",
      "358249491\n",
      "36279\n",
      "363655532\n",
      "366406728\n",
      "390959699\n",
      "391044439\n",
      "406403918\n",
      "414380191\n",
      "440738016\n",
      "441203627\n",
      "442692020\n",
      "453009027\n",
      "454140642\n",
      "1002795171\n",
      "274627577\n",
      "441721286\n",
      "88992642\n",
      "99784331\n",
      "259169431\n",
      "339053196\n",
      "352140271\n",
      "404842248\n",
      "302965287\n",
      "389124828\n",
      "384490098\n",
      "272466346\n",
      "246290171\n",
      "263015893\n",
      "33862068\n",
      "33877338\n",
      "131857479\n",
      "13346068\n",
      "146579304\n",
      "153408739\n",
      "174203886\n",
      "327681256\n",
      "387836690\n",
      "199523754\n",
      "17131560\n",
      "118182747\n",
      "144401325\n",
      "641539146\n",
      "374137958\n",
      "148963635\n",
      "250496108\n",
      "2008877991\n",
      "160600809\n",
      "292464126\n",
      "165428216\n",
      "3311888\n",
      "356630364\n",
      "233714249\n",
      "230386284\n",
      "326919774\n",
      "445706410\n",
      "287135515\n",
      "182650117\n",
      "280097876\n",
      "394358240\n",
      "119087165\n",
      "6352545\n",
      "50308490\n",
      "326903459\n",
      "370134610\n",
      "41111279\n",
      "238180562\n",
      "293571051\n",
      "364425591\n",
      "306196907\n",
      "167724589\n",
      "253918451\n",
      "375090717\n",
      "456588598\n",
      "229434615\n",
      "430689771\n",
      "252653577\n",
      "188766090\n",
      "229455001\n",
      "349154278\n",
      "339735718\n",
      "364873540\n",
      "279956763\n",
      "294393328\n",
      "32000208\n",
      "90563464\n",
      "150358433\n",
      "303910522\n",
      "167792839\n",
      "290605311\n",
      "279957439\n",
      "279955276\n",
      "75683740\n",
      "90562888\n",
      "88129419\n",
      "9779230\n",
      "299156852\n",
      "219485088\n",
      "391081821\n",
      "185011414\n",
      "246288840\n",
      "348899878\n",
      "30906987\n",
      "425083216\n",
      "404799340\n",
      "304864057\n",
      "306270678\n",
      "373840212\n",
      "89740313\n",
      "230966538\n",
      "464488581\n",
      "300734224\n",
      "267434733\n",
      "287833838\n",
      "424879951\n",
      "272358852\n",
      "114271331\n",
      "411829582\n",
      "119384992\n",
      "91001369\n",
      "150394361\n",
      "364425576\n",
      "381291287\n",
      "391821647\n",
      "120598737\n",
      "26821153\n",
      "258514380\n",
      "156458268\n",
      "163965682\n",
      "390349501\n",
      "640919535\n",
      "40993256\n",
      "436932094\n",
      "436005037\n",
      "42321328\n",
      "309682332\n",
      "428331598\n",
      "273091578\n",
      "94368545\n",
      "148574197\n",
      "289228040\n",
      "154091615\n",
      "252365657\n",
      "381202615\n",
      "70013723\n",
      "99396894\n",
      "23079563\n",
      "467354611\n",
      "255991074\n",
      "364152478\n",
      "369834384\n",
      "273417145\n",
      "260986765\n",
      "25078924\n",
      "61614242\n",
      "183981567\n",
      "272316843\n",
      "30193373\n",
      "277026183\n",
      "128846212\n",
      "61136885\n",
      "177238006\n",
      "49449799\n",
      "379192021\n",
      "301896004\n",
      "27968643\n",
      "101654574\n",
      "140691839\n",
      "177389419\n",
      "339633666\n",
      "274017873\n",
      "306108833\n",
      "277027632\n",
      "177352566\n",
      "404933877\n",
      "373843852\n",
      "99222716\n",
      "332051907\n",
      "841054407\n",
      "286945730\n",
      "166052265\n",
      "354785553\n",
      "159173586\n",
      "117631084\n",
      "159840569\n",
      "259879138\n",
      "236896126\n",
      "414897731\n",
      "392426715\n",
      "354959477\n",
      "664035272\n",
      "302698268\n",
      "174042805\n",
      "229751802\n",
      "236788144\n",
      "441804674\n",
      "377197068\n",
      "298295539\n",
      "18252911\n",
      "61110114\n",
      "139258713\n",
      "120602501\n",
      "260593507\n",
      "367017335\n",
      "9483535\n",
      "117247507\n",
      "298896016\n",
      "312666378\n",
      "270217681\n",
      "305479069\n",
      "27434899\n",
      "268308889\n",
      "277316602\n",
      "281667628\n",
      "300029336\n",
      "303557843\n",
      "305582965\n",
      "307999339\n",
      "309472493\n",
      "20925093\n",
      "209407002\n",
      "219332212\n",
      "220474216\n",
      "234331809\n",
      "237944314\n",
      "248988133\n",
      "260213488\n",
      "356744650\n",
      "42837303\n",
      "331387944\n",
      "20352516\n",
      "229502230\n",
      "393952062\n"
     ]
    }
   ],
   "source": [
    "for i in pl[neg & ppos]['TIC ID']:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "index 0 is out of bounds for axis 0 with size 0",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-c6cfa0952c23>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0mcompare\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mensemble_preds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpreds\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mpl\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpl\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'TIC ID'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m118412801\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m: index 0 is out of bounds for axis 0 with size 0"
     ]
    }
   ],
   "source": [
    "def compare(ensemble_preds, filter):\n",
    "    result = ensemble_preds[0][filter]\n",
    "    for preds in ensemble_preds[1:]:\n",
    "        result = result.append(preds[filter])\n",
    "    return result\n",
    "\n",
    "compare(ensemble_preds, preds.index == pl[pl['TIC ID'] == 118412801].index.values[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl[pl.index == pl[pl['TIC ID'] == 1254504863].index.values[0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PR curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids = set(ensemble_preds[0].index.values)\n",
    "\n",
    "index = {v: i for i, v in enumerate(ids)}\n",
    "\n",
    "pred_es = np.zeros([len(ensemble_preds), len(index)])\n",
    "for i, preds in enumerate(ensemble_preds):\n",
    "    for row in preds.iterrows():\n",
    "        ex_id, pred_e = row[0], row[1][col_e]\n",
    "        pred_es[i][index[ex_id]] = pred_e\n",
    "\n",
    "lbl_es = np.zeros([len(index)], dtype=np.bool)\n",
    "for row in tce_labels.iterrows():\n",
    "    ex_id, lbl_e = row[0], row[1]['disp_e']\n",
    "    lbl_es[index[ex_id]] = (lbl_e > 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cond_pos = int(np.sum(lbl_es))\n",
    "\n",
    "def pr_at_th(th):\n",
    "    pred_pos = np.any(pred_es >= th, axis=0)\n",
    "    true_pos = pred_pos & lbl_es\n",
    "    num_pred_pos = int(np.sum(pred_pos))\n",
    "    num_true_pos = int(np.sum(true_pos))\n",
    "    if num_pred_pos == 0:\n",
    "        return 1.0, 0.0\n",
    "    return float(num_true_pos) / float(num_pred_pos), float(num_true_pos) / float(num_cond_pos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "ps, rs, ths = ([], [], [])\n",
    "th = np.max(pred_es)\n",
    "while th >= 0.0:\n",
    "    p, r = pr_at_th(th)\n",
    "    ps.append(p)\n",
    "    rs.append(r)\n",
    "    ths.append(th)\n",
    "    th -= 0.0005\n",
    "    \n",
    "from sklearn import metrics\n",
    "\n",
    "print(f'AUC: {metrics.auc(rs, ps)}, max R: {max(rs)}, max P: {max(ps)}')\n",
    "    \n",
    "i = len(rs) - 1\n",
    "while rs[i] == 1.0:\n",
    "    i -= 1\n",
    "i += 1\n",
    "print(f'100% recall at: {int(ps[i] * 100)}%, threshold: {ths[i]}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 3.7), dpi=200)\n",
    "\n",
    "ax.spines['top'].set_color('#808080')\n",
    "ax.spines['right'].set_color('#808080')\n",
    "ax.spines['left'].set_color('#808080')\n",
    "ax.spines['bottom'].set_color('#808080')\n",
    "ax.tick_params(direction='in', color='#808080')\n",
    "\n",
    "plt.grid(color='#c0c0c0', linestyle='--', linewidth=0.5)\n",
    "\n",
    "plt.ylabel('Precision', fontweight='bold')\n",
    "plt.xlabel('Recall', fontweight='bold')\n",
    "\n",
    "plt.xlim((0.0, 1.0))\n",
    "plt.ylim((0.0, 1.0))\n",
    "\n",
    "_ = plt.plot(rs, ps)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
